//
// Created by ljy on 19-4-22.
//

#include <fstream>
#include "complx_tracker.h"

#define DEBUG_SHOW_IMG
#define OUTPUT_ARTICLE_IMAGE false

using namespace std;
using namespace cv;

pointsDetector::pointsDetector() {
    if (algo == "MOG2")
        pBackSub = createBackgroundSubtractorMOG2(500, 16, false); //500 400 false
    else if (algo == "KNN")
        pBackSub = cv::bgsegm::createBackgroundSubtractorCNT(300, 400, false);
    else if (algo == "MOG")
        pBackSub = bgsegm::createBackgroundSubtractorMOG(500, 10, 0.7, 0);

    imgRec.push_back(Mat::zeros(img.size(),img.type()));
    imgRec.push_back(Mat::zeros(img.size(),img.type()));
    imgRec.push_back(Mat::zeros(img.size(),img.type()));
    imgRec.push_back(Mat::zeros(img.size(),img.type()));
    imgRec.push_back(Mat::zeros(img.size(),img.type()));

}


bool pointsDetector::updateFrame(const cv::Mat &frame, bool using_points_position_guess) {
    if (frame.empty()) {
        ROS_INFO("pointsDetector: Can't get frame.");
        return false;
    }
    img = frame;
    valueImg = Mat::zeros(img.size(), CV_8UC1);
    cvtColor(frame, hsvImg, CV_BGR2HSV);

    Mat mask;
    inRange(hsvImg, Scalar(0, 200, 100), Scalar(10, 255, 255), mask);
//    hsvImg.copyTo(hsvImg,mask);
    for (int r = 0; r < hsvImg.rows; r++)
    {
        for (int c = 0; c < hsvImg.cols; c++)
        {
            if (mask.at<uchar>(r, c) == 255)
            {
                valueImg.at<uchar>(r, c) = hsvImg.at<Vec3b>(r, c)[2];
            }
        }
    }
//
//    for (int i = 0; i < img.rows; i++) {
//        for (int j = 0; j < img.cols; j++) {
//            const int hi = i * img.cols * 3 + j * 3,
//                    gi = i * img.cols + j;
//            valueImg.data[gi] = hsvImg.data[hi + 2];
//        }
//    }


    if(OUTPUT_ARTICLE_IMAGE)
    {
        imgRec[0] = img;
        imgRec[1] = valueImg;
    }

    if(!using_points_position_guess)
    {
        //! [apply]
        //update the background model
        pBackSub->apply(img, fgMask);
        //! [apply]
        frame_masked = Mat::zeros(frame.size(), frame.type());
        valueImg.copyTo(frame_masked, fgMask);
        if(OUTPUT_ARTICLE_IMAGE) {
            imgRec[2] = frame_masked;
        }

    } else{


        frame_masked = Mat::zeros(frame.size(), frame.type());
        valueImg.copyTo(frame_masked);
        if(OUTPUT_ARTICLE_IMAGE) {
            //!for imgRec
            //! [apply]
            //update the background model
            pBackSub->apply(img, fgMask);
            //! [apply]
            imgRec[2] = Mat::zeros(valueImg.size(), valueImg.type());
            valueImg.copyTo(imgRec[2], fgMask);
        }
    }
}


bool pointsDetector::getPoints(const Mat &frame, std::vector<cv::Point2f> &outputPointsVector,
                               bool using_points_position_guess) {

    if(!using_points_position_guess)
    {
        updateFrame(frame,false);
    }

    count++;


    threshold(frame_masked, binImg, 10, 255, THRESH_BINARY);

//    imshow("frame_masked", frame_masked);

    ////dilate
    Mat dilatedImg;
    int structElementSize = 1;
    Mat structElement = getStructuringElement(MORPH_ELLIPSE,
                                              Size(2 * structElementSize + 1, 2 * structElementSize + 1),
                                              Point(structElementSize, structElementSize));
    dilate(binImg, dilatedImg, structElement);

    ////find contours
    vector<vector<Point>> contours, contours_org;

    vector<Vec4i> hierarchy;
    findContours(dilatedImg, contours_org, hierarchy, RETR_CCOMP, CHAIN_APPROX_NONE);



    contoursImg = Mat::zeros(img.size(), CV_8UC1);

    vector<myImgPoint> imgPts;


    for (auto itContour = contours_org.begin(); itContour != contours_org.end(); itContour++) {
//            if(contourArea(*itContour)<40)
//            {
//                continue;
//            }

        Moments m;
        m = moments(*itContour, true);
        myImgPoint p = myImgPoint((m.m10 / m.m00), (m.m01 / m.m00));
        imgPts.push_back(p);

        for (auto itContourPt = (*itContour).begin(); itContourPt != (*itContour).end(); itContourPt++) {
            contoursImg.at<uchar>(Point((*itContourPt).x, (*itContourPt).y)) = 255;
        }
        contoursImg.at<uchar>(p) = 255;
    }

    if(OUTPUT_ARTICLE_IMAGE) {
        imgRec[3] = contoursImg;
    }
//        contoursImg.at<uchar>(Point2f(50,200))=255;
//    imshow("contours", contoursImg);

    if (using_points_position_guess) {
        for (int i = 0; i < 4; i++) {
            float min_dis = 9999;
            auto closest_pt = imgPts.begin();
            for (auto itPt = imgPts.begin(); itPt != imgPts.end(); itPt++) {
                float new_dis = myImgPoint::dist_btn_pts(*itPt, outputPointsVector[i]);
                if (new_dis < min_dis) {
                    min_dis = new_dis;
                    closest_pt = itPt;
                }
            }
            if (min_dis > 100) {
                using_points_position_guess = false;
                break;
//                    return false;
            }
            sortedImgPts[i] = *closest_pt;
        }
    }
    if (using_points_position_guess) {
        for (int i = 0; i < 4; i++) {
            for (int j = i + 1; j < 4; j++) {
                if (myImgPoint::dist_btn_pts(sortedImgPts[i], sortedImgPts[j]) < 5) {
                    using_points_position_guess = false;
                    break;
//                        return false;

                }
            }
        }
    }

    if (!using_points_position_guess) {

        //region find max_Area
        double maxArea;
        vector<Point> temp;
        for (auto i = contours_org.begin(); i != contours_org.end(); i++) {
            maxArea = contourArea(*i);
            for (auto j = i + 1; j != contours_org.end(); j++) {
                if (contourArea(*j) > maxArea) {
                    maxArea = contourArea(*j);
                    temp = *i;
                    (*i) = (*j);
                    (*j) = temp;
                }
            }
        }
        if (contours_org.size() >= 4)
            for (int i = 0; i < 4; i++) contours.push_back(contours_org[i]);
        //endregion

        contoursImg = Mat::zeros(img.size(), CV_8UC1);

        vector<myImgPoint> imgPts;


        for (auto itContour = contours.begin(); itContour != contours.end(); itContour++) {
//            if(contourArea(*itContour)<40)
//            {
//                continue;
//            }

            Moments m;
            m = moments(*itContour, true);
            myImgPoint p = myImgPoint((m.m10 / m.m00), (m.m01 / m.m00));
            imgPts.push_back(p);

            for (auto itContourPt = (*itContour).begin(); itContourPt != (*itContour).end(); itContourPt++) {
                contoursImg.at<uchar>(Point((*itContourPt).x, (*itContourPt).y)) = 255;
            }
            contoursImg.at<uchar>(p) = 255;
        }


        if (imgPts.size() == 4 && count >= 10) {


            Point2f centerP;
            sortedImgPts.clear();
            sortedImgPts.assign(4, myImgPoint(0, 0));
            centerP = Point2f(((imgPts[0].x + imgPts[1].x + imgPts[2].x + imgPts[3].x) / 4.0),
                              ((imgPts[0].y + imgPts[1].y + imgPts[2].y + imgPts[3].y) / 4.0));

            if (DETECTING_METHOD == 0) {
                //与中心点形成的角度
                for (int i = 0; i < 4; i++) {
                    imgPts[i].angle_with_center = atan2(imgPts[i].y - centerP.y, imgPts[i].x - centerP.x);
                }
                //sort according to angle with center Point
                sort(imgPts.begin(), imgPts.end());
                imgPts[0].length_to_next_point = myImgPoint::dist_btn_pts(imgPts[0], imgPts[1]);
                imgPts[1].length_to_next_point = myImgPoint::dist_btn_pts(imgPts[1], imgPts[2]);
                imgPts[2].length_to_next_point = myImgPoint::dist_btn_pts(imgPts[2], imgPts[3]);
                imgPts[3].length_to_next_point = myImgPoint::dist_btn_pts(imgPts[3], imgPts[0]);
                int Index_maxL = 0;
                for (int i = 1; i < 4; i++) {
                    if (imgPts[i].length_to_next_point > imgPts[Index_maxL].length_to_next_point)
                        Index_maxL = i;
                }
                sortedImgPts[0] = imgPts[Index_maxL];
                sortedImgPts[1] = imgPts[(Index_maxL + 1) % 4];
                sortedImgPts[2] = imgPts[(Index_maxL + 2) % 4];
                sortedImgPts[3] = imgPts[(Index_maxL + 3) % 4];
            } else if (DETECTING_METHOD == 1) {
                for (int i = 0; i < 4; i++) {
                    if (imgPts[i].x < centerP.x && imgPts[i].y < centerP.y) {
                        sortedImgPts[3] = imgPts[i];
                    } else if (imgPts[i].x < centerP.x && imgPts[i].y >= centerP.y) {
                        sortedImgPts[0] = imgPts[i];
                    } else if (imgPts[i].x >= centerP.x && imgPts[i].y >= centerP.y) {
                        sortedImgPts[1] = imgPts[i];
                    } else if (imgPts[i].x >= centerP.x && imgPts[i].y < centerP.y) {
                        sortedImgPts[2] = imgPts[i];
                    }
                }
            }else if (DETECTING_METHOD == 2) {
                int up = 99999, down = 0, left = 99999, right = 0;
                int iup,idown,ileft,iright;
                for (int i = 0; i < 4; i++) {
                    if (imgPts[i].y < up) {
                        up = imgPts[i].y;
                        iup = i;
                    }
                    if (imgPts[i].y > down) {
                        down = imgPts[i].y;
                        idown = i;
                    }
                    if (imgPts[i].x < left) {
                        left = imgPts[i].x;
                        ileft = i;
                    }
                    if (imgPts[i].x > right) {
                        right = imgPts[i].x;
                        iright = i;
                    }
                }
                sortedImgPts[0] = imgPts[iup];
                sortedImgPts[1] = imgPts[idown];
                sortedImgPts[2] = imgPts[ileft];
                sortedImgPts[3] = imgPts[iright];
            }
        } else {
            return false;
        }
    }
    outputPointsVector = sortedImgPts;
    if(OUTPUT_ARTICLE_IMAGE) {
        img.copyTo(imgRec[4]);
        putText(imgRec[4], "0", sortedImgPts[0], FONT_HERSHEY_PLAIN, 2, Scalar(0, 0, 255));
        putText(imgRec[4], "1", sortedImgPts[1], FONT_HERSHEY_PLAIN, 1, Scalar(0, 0, 255));
        putText(imgRec[4], "2", sortedImgPts[2], FONT_HERSHEY_PLAIN, 1, Scalar(0, 0, 255));
        putText(imgRec[4], "3", sortedImgPts[3], FONT_HERSHEY_PLAIN, 1, Scalar(0, 0, 255));
    }
    return true;

}

bool pointsFollower::getPoints(const cv::Mat &frame, std::vector<cv::Point2f> &outputPointsVector) {

    if (frame.empty()) {
        ROS_INFO("pointsFollower: Can't get frame.");
        return false;
    }
    Mat img, hsvImg;
    gray = Mat::zeros(frame.size(), CV_8UC1);
    img = frame;
    cvtColor(frame, hsvImg, CV_BGR2HSV);
    for (int i = 0; i < img.rows; i++) {
        for (int j = 0; j < img.cols; j++) {
            const int hi = i * img.cols * 3 + j * 3,
                    gi = i * img.cols + j;
            gray.data[gi] = hsvImg.data[hi + 2];
        }
    }
    prevImgPts = outputPointsVector;
    nextImgPts.resize(prevImgPts.size());

    vector<uchar> status;
    vector<float> err;
//    TermCriteria termcrit(TermCriteria::COUNT|TermCriteria::EPS,50,0.001);
    TermCriteria termcrit(TermCriteria::EPS, 50, 0.001);
    Size winSize(31, 31);


    if (outputPointsVector.size() != 4) {
        swap(gray, prevGray);
        ROS_INFO("pointsFollower: Number of Points to be followed is wrong.");
        return false;
    }
    if (prevGray.empty())
        gray.copyTo(prevGray);
//    calcOpticalFlowPyrLK(prevGray, gray, prevImgPts, nextImgPts, status, err);
    calcOpticalFlowPyrLK(prevGray, gray, prevImgPts, nextImgPts, status, err, winSize,
                         0, termcrit, 0, 0.001);
    size_t i, k;
    for (i = k = 0; i < nextImgPts.size(); i++) {
        if (!status[i])
            continue;
        nextImgPts[k++] = nextImgPts[i];
//        circle( depth_img, nextImgPts[i], 2, Scalar(0,255,0), -1, 8);
    }
    if (prevImgPts.size() == nextImgPts.size()) {
        outputPointsVector = nextImgPts;
        swap(gray, prevGray);
        swap(prevImgPts, nextImgPts);
        return true;
    } else {
        swap(gray, prevGray);
        ROS_INFO("pointsFollower: failed.");
        return false;
    }
}

complx_tracker::complx_tracker() {
    cameraMatrix = cv::Mat::zeros(3, 3, CV_64F);
    distCoeffs = std::vector<double>{0, 0, 0,
                                     0, 0};
    cameraMatrix.at<double>(0, 0) = 277.191356*1.5;
    cameraMatrix.at<double>(0, 2) = 320.5*1.5;
    cameraMatrix.at<double>(1, 1) = 277.191356*1.5;
    cameraMatrix.at<double>(1, 2) = 240.5*1.5;
    cameraMatrix.at<double>(2, 2) = 1;

    start_time = ros::Time::now();

    if (STRUCTURE_INDEX == 0) {
        objectPts.emplace_back(-0.071, -0.054, 0);
        objectPts.emplace_back(0.071, -0.054, 0);
        objectPts.emplace_back(0.029, 0.054, 0);
        objectPts.emplace_back(-0.029, 0.053, 0);
    } else if (STRUCTURE_INDEX == 1) {
        objectPts.emplace_back(0.01, 0.0925, 0);
        objectPts.emplace_back(0.01, -0.0925, 0);
        objectPts.emplace_back(0.01, -0.078, 0.088);
        objectPts.emplace_back(0.01, 0.078, 0.088);
    } else if (STRUCTURE_INDEX == 2) {
        objectPts.emplace_back(0, 0.098, 0 - 0.01);
        objectPts.emplace_back(0, -0.098, 0 - 0.01);
        objectPts.emplace_back(0, -0.076, 0.086 - 0.01);
        objectPts.emplace_back(0, 0.076, 0.086 - 0.01);
    } else if (STRUCTURE_INDEX == 3) {
        objectPts.emplace_back(0, 0,  0.35);
        objectPts.emplace_back(0, 0, -0.35);
        objectPts.emplace_back(0, 0.35, 0);
        objectPts.emplace_back(0, -0.35, 0);
    }


    imgPts = std::vector<cv::Point2f>(4);

    std::string pkg_path = ros::package::getPath("drone_pnp");


    Matx44d trans_cam;
    ifstream mtx;
    mtx.open(pkg_path + "/data/transformation_matrix.txt");
    if (mtx.is_open()) {
        mtx >> trans_cam.val[0] >> trans_cam.val[1] >> trans_cam.val[2] >> trans_cam.val[3]
            >> trans_cam.val[4] >> trans_cam.val[5] >> trans_cam.val[6] >> trans_cam.val[7]
            >> trans_cam.val[8] >> trans_cam.val[9] >> trans_cam.val[10] >> trans_cam.val[11]
            >> trans_cam.val[12] >> trans_cam.val[13] >> trans_cam.val[14] >> trans_cam.val[15];
        mtx.close();
    } else
        ROS_ERROR("Can't read matrix file.");
    camPose = Affine3d(trans_cam);

    following_flag = false;
    myWindow = viz::Viz3d("Coordinate Frame");
    myWindow.showWidget("Coordinate Widget", viz::WCoordinateSystem());
    cube_widget1 = viz::WCube(Point3f(-0.0663, -0.03966, -0.01), Point3f(0.0663, 0, 0.01), false, viz::Color::blue());
    cube_widget2 = viz::WCube(Point3f(-0.03135, 0, -0.01), Point3f(0.03135, 0.06132, 0.01), false, viz::Color::blue());
    myWindow.showWidget("Cube Widget1", cube_widget1);
    myWindow.showWidget("Cube Widget2", cube_widget2);

    cv::viz::WLine line_widget(Point3d(0, 0, 3), Point3d(0, 0, 0), viz::Color::red());
    myWindow.showWidget("line", line_widget);
    line_widget.setRenderingProperty(viz::LINE_WIDTH, 1);


}


bool complx_tracker::findImgPts() {
    if (!following_flag) {
        if (pts_detector.getPoints(img, imgPts)) {
            following_flag = true;
        } else
            return false;
    } else {
        following_flag = pts_follower.getPoints(img, imgPts);
        pts_detector.updateFrame(img,true);
        for (int i = 0; i < 4; i++) {
            int t = pts_detector.valueImg.at<uchar>(imgPts[i]);
            if (t < 50) {
                following_flag = pts_detector.getPoints(img, imgPts, true);

                if(OUTPUT_ARTICLE_IMAGE) {
                    std::string pkg_path = ros::package::getPath("drone_pnp");
                    int rand_int = random() % 100;
                    for (int saveImgIndex = 0; saveImgIndex < pts_detector.imgRec.size(); saveImgIndex++) {
                        imwrite(pkg_path + "/image/" + to_string(rand_int) + to_string(saveImgIndex) + ".jpg",
                                pts_detector.imgRec[saveImgIndex]);
                    }
                }

                break;
            }
        }
//        if(!following_flag)
//            return findImgPts();
    }
    return following_flag;
}

bool complx_tracker::apply(const cv::Mat &frame, cv::Mat &img_to_show, cv::Affine3d &pose_world, cv::Affine3d &pose_raw) {
    //undistort(frame, img, cameraMatrix, distCoeffs);
    frame.copyTo(img);


    findImgPts();

    solvePnP(objectPts, imgPts, cameraMatrix, distCoeffs, outputRvecRaw, outputTvecRaw);
    pose = Affine3d(outputRvecRaw, outputTvecRaw);
    pose = pose.inv();// from drone to ground
    worldPose = pose.concatenate(camPose);
    pose_world = worldPose;
    pose_raw = pose;

    virtualize();
    img_to_show = showImg;
    return true;

}


void complx_tracker::virtualize() {
//    myWindow.setWidgetPose("Cube Widget1", worldPose);
//    myWindow.setWidgetPose("Cube Widget2", worldPose);
//    myWindow.spinOnce(1, true);
    img.copyTo(showImg);
    if (!pts_follower.nextImgPts.empty()) {
//        circle( showImg, pts_follower.nextImgPts[0], 1, Scalar(0,255,0), -1, 8);
//        circle( showImg, pts_follower.nextImgPts[1], 1, Scalar(255,255,0), -1, 8);
//        circle( showImg, pts_follower.nextImgPts[2], 1, Scalar(0,255,255), -1, 8);
//        circle( showImg, pts_follower.nextImgPts[3], 1, Scalar(255,0,255), -1, 8);

    }
//    imshow("contours",pts_detector.contoursImg);
    if (!pts_detector.sortedImgPts.empty()) {
//        circle( showImg, pts_detector.sortedImgPts[0], 3, Scalar(255,0,0), -1, 8);
//        circle( showImg, pts_detector.sortedImgPts[1], 3, Scalar(255,0,0), -1, 8);
//        circle( showImg, pts_detector.sortedImgPts[2], 3, Scalar(255,0,0), -1, 8);
//        circle( showImg, pts_detector.sortedImgPts[3], 3, Scalar(255,0,0), -1, 8);

    }
    if (!imgPts.empty()) {
//        circle( showImg, imgPts[0], 3, Scalar(0,255,0), -1, 8);
//        circle( showImg, imgPts[1], 3, Scalar(255,255,0), -1, 8);
//        circle( showImg, imgPts[2], 3, Scalar(0,255,255), -1, 8);
//        circle( showImg, imgPts[3], 3, Scalar(255,0,255), -1, 8);

        putText(showImg, "0", imgPts[0], FONT_HERSHEY_PLAIN, 2, Scalar(0, 255, 0));
        putText(showImg, "1", imgPts[1], FONT_HERSHEY_PLAIN, 1, Scalar(0, 255, 0));
        putText(showImg, "2", imgPts[2], FONT_HERSHEY_PLAIN, 1, Scalar(0, 255, 0));
        putText(showImg, "3", imgPts[3], FONT_HERSHEY_PLAIN, 1, Scalar(0, 255, 0));
    }

    imshow("tracking", showImg);
    waitKey(1);
}





















